#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
# this is copied from df/outcore.py
#


import tempfile
import numpy
import scipy.linalg
import h5py
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.outcore import _load_from_h5g
from pyscf import __config__

MAX_MEMORY = getattr(__config__, 'df_outcore_max_memory', 2000)  # 2GB
LINEAR_DEP_THR = getattr(__config__, 'df_df_DF_lindep', 1e-12)

def general(mol, mo_coeffs, idxrange, erifile, dataname='eri_mo', tmpdata='eri_mo', tmperi='aocderi.dat',
            aosym='s2ij', comp=1, max_memory=MAX_MEMORY, verbose=0, compact=True):
    ''' Transform ij of (ij|L) saved in temperi to MOs.
        This is copied from general with cholesky_eri_b removed
    '''
    assert(aosym in ('s1', 's2ij'))
    time0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mol, verbose)

    fswap = h5py.File(tmperi, 'r')
    time1 = log.timer('AO->MO eri transformation 1 pass', *time0)

    nao = mo_coeffs[0].shape[0]
    if aosym == 's1':
        nao_pair = nao * nao
        aosym_as_nr_e2 = 's1'
    else:
        nao_pair = nao * (nao+1) // 2
        aosym_as_nr_e2 = 's2kl'

    ijmosym, nij_pair, moij, ijshape = \
            ao2mo.incore._conc_mos(mo_coeffs[0], mo_coeffs[1],
                                   compact and aosym != 's1')

    naoaux = fswap['%s/0'%tmpdata].shape[-2]
    ndata = idxrange[1] - idxrange[0]
    feri = _create_h5file(erifile, dataname)
    if comp == 1:
        h5d_eri = feri.create_dataset(dataname, (ndata,nij_pair), 'f8')
    else:
        h5d_eri = feri.create_dataset(dataname, (comp,ndata,nij_pair), 'f8')

    def load(row_slice):
        row0, row1 = row_slice
        row0 += idxrange[0]
        row1 += idxrange[0]
        return _load_from_h5g(fswap[tmpdata], row0, row1)

    iolen = min(max(int(max_memory*.45e6/8/(nao_pair+nij_pair)), 28), ndata)
    totstep = (ndata+iolen-1)//iolen
    ti0 = time1
    slices = list(lib.prange(0, ndata, iolen))
    for istep, dat in enumerate(lib.map_with_prefetch(load, slices)):
        row0, row1 = slices[istep]
        nrow = row1 - row0
        if comp == 1:
            dat = _ao2mo.nr_e2(dat, moij, ijshape, aosym_as_nr_e2, ijmosym)
            h5d_eri[row0:row1] = dat
        else:
            dat = _ao2mo.nr_e2(dat.reshape(comp*nrow, nao_pair),
                               moij, ijshape, aosym_as_nr_e2, ijmosym)
            h5d_eri[:,row0:row1] = dat.reshape(comp, nrow, nij_pair)
        dat = None
        log.debug('step 2 [%d/%d], [%d:%d], row = %d',
                  istep+1, totstep, row0, row1, nrow)
        ti0 = log.timer('step 2 [%d/%d], [%d:%d], row = %d'%
                        (istep+1, totstep, row0, row1, nrow), *ti0)

    fswap.close()
    feri.close()
    log.timer('AO->MO CD eri transformation 2 pass', *time1)
    log.timer('AO->MO CD eri transformation', *time0)
#   return erifile
    return ndata

def _create_h5file(erifile, dataname):
    if h5py.is_hdf5(erifile):
        feri = h5py.File(erifile, 'a')
        if dataname in feri:
            del(feri[dataname])
    else:
        feri = h5py.File(erifile, 'w')
    return feri

del(MAX_MEMORY)


